# --------------------------
# -*- coding=utf-8 -*-
#
# written by zhaoxingjie
# 18829350080@163.com
# --------------------------
import glob
import json
import os
import os.path as osp
from xml.etree import ElementTree as ET

import cv2
from pycocotools.coco import COCO

# test
# cocojsn = '/media/hp208/4t/data/coco/2017/annotations/instances_val2017.json'
#
# with open(cocojsn, 'r') as f:
#     res = json.load(f)

"""
[
    [
        {},
        {}
    ],
    [
        {}
        {}
    ]
]

licenses
    url
images
    id
    date_captured
    license
    file_name
    coco_url
    flickr_url
    width
    height
categories
    id
    supercategory
    name
annotations
    id
    bbox
    iscrowd
    segmentation
    category_id
    area
    image_id
info
    verson
    year
    describtion
"""
licenses = [dict(url='https://www.gitee.com/zhaowujie')]
info = [dict(version=0.1,
             year=2020,
             describtion='Anti-UAV2020 1st',
             )]
categories = [dict(id=1, supercategory='uav', name='uav')]

image_ids = 0
annotations_ids = 0
images = []
annotations = []

cat2idx = {'uav': 1}
image_root = '/media/hp208/4t/zhaoxingjie/project_graduation/data/test-dev-voc-format-ir/VOC2007/JPEGImages/'
xml_root = '/media/hp208/4t/zhaoxingjie/project_graduation/data/test-dev-voc-format-ir/VOC2007/Annotations/*'

img_paths = glob.glob(image_root)
xml_paths = glob.glob(xml_root)


def get_set_paths(set_root, txt_path):
    # paths = []
    with open(txt_path, 'r') as f:
        lines = f.readlines()
    paths = [osp.join(set_root, line.strip() + '.jpg') for line in lines]

    return paths


def convert(img_paths, json_name):
    """
    convert .xml to .json
    :param img_paths: image_paths, list
    :param json_name: file to write
    :return: None
    """
    global image_ids
    global annotations_ids
    for idx, image_path in enumerate(img_paths):
        assert os.path.isfile(image_path), "File does not exist! "
        xml_path = image_path.replace('JPEGImages', 'Annotations').replace('jpg', 'xml')
        # ----- find image info ----- #
        img_dict = dict(id=image_ids)
        img_file_neme = osp.basename(image_path)
        img_dict['file_name'] = img_file_neme
        height, width, _ = cv2.imread(image_path).shape
        img_dict['width'] = width
        img_dict['height'] = height
        images.append(img_dict)

        # parse .xml
        with open(xml_path, 'r') as f:
            root = ET.parse(f).getroot()
        for obj in root.findall('object'):
            cls_name = obj.find('name').text
            bndbox = obj.find('bndbox')
            xmin = int(bndbox.find('xmin').text)
            ymin = int(bndbox.find('ymin').text)
            xmax = int(bndbox.find('xmax').text)
            ymax = int(bndbox.find('ymax').text)
            # ----- find anno info ----- #
            anno_dict = dict(id=annotations_ids)
            anno_dict['segmentation'] = []
            anno_dict['area'] = (xmax - xmin) * (ymax - ymin)
            anno_dict['image_id'] = image_ids
            anno_dict['category_id'] = cat2idx[cls_name]
            anno_dict['iscrowd'] = 0
            anno_dict['bbox'] = [xmin, ymin, xmax - xmin, ymax - ymin]
            annotations.append(anno_dict)
            annotations_ids += 1  # next object

        image_ids += 1  # next img and anno
        print('Converting: >>> {}/{}'.format(idx, len(img_paths)), flush=True)
    # write to dict
    result = dict(licenses=licenses,
                  images=images,
                  annotations=annotations,
                  info=info,
                  categories=categories,
                  )
    with open(json_name, 'w') as f:
        json.dump(result, f)
        print('Save json file to {}'.format(json_name))


img_set_root = '/media/hp208/4t/zhaoxingjie/project_graduation/data/test-dev-voc-format-ir/VOC2007/ImageSets/Main'
train_txt = osp.join(img_set_root, 'trainval.txt')
test_txt = osp.join(img_set_root, 'test.txt')
val_txt = osp.join(img_set_root, 'val.txt')

paths = get_set_paths(image_root, val_txt)
convert(paths, 'val.json')

# -------TEST ------- #
coco = COCO(annotation_file='val.json')
print("coco\nimages.size [%05d]\tannotations.size [%05d]\t category.size [%05d]\ndone!"
      % (len(coco.imgs), len(coco.anns), len(coco.cats)))
